import base64
import json
from typing import List

from volcengine.visual.VisualService import VisualService
from ..__result import OperateResult


class VolcengineClient(VisualService):
    def __init__(self, volcengine_info=None, use_url: bool = True):
        super().__init__()
        volcengine_info = volcengine_info or {}
        if not volcengine_info:
            raise Exception("No VolcengineInfo found")
        if isinstance(volcengine_info, str):
            volcengine_info = json.loads(volcengine_info)
        self.volcengine_info = volcengine_info
        self.set_ak(self.volcengine_info.get("ak", ""))
        self.set_sk(self.volcengine_info.get("sk", ""))
        self.set_socket_timeout(30)
        self.set_connection_timeout(30)
        self.set_connection_timeout(30)
        self.set_socket_timeout(30)
        self.use_url = use_url

    def __download(self, url):
        response = self.session.get(url)
        return response.content

    def __common_handler_response(self, response: dict = None) -> OperateResult:
        if isinstance(response, dict):
            code = int(response.get("code", 400))
            if code == 10000:
                image_data = response.get("data", {})
                images = []
                if not self.use_url:
                    base64_images = image_data.get("binary_data_base64", [])
                    for image_index, base64_image in enumerate(base64_images):
                        image_bytes = base64.b64decode(base64_image)
                        images.append(image_bytes)
                else:
                    image_urls = image_data.get("image_urls", [])
                    for image_index, image_url in enumerate(image_urls):
                        image_bytes = self.__download(image_url)
                        images.append(image_bytes)
                if images:
                    return OperateResult.success(data=images)
                return OperateResult.warning(message="生成失败")
            else:
                return OperateResult.warning(message=response.get("message", "任务失败"))
        else:
            return OperateResult.warning(message=str(response))

    def seed_edit(self, image_url: str = None, image_base64: str = None, prompt: str = "图片细节优化一下",
                  negative_prompt: str = "Low quality", scale: float = 0.5) -> OperateResult:
        form = {
            "req_key": "byteedit_v2.0",
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "seed": -1,
            "scale": scale,
            "return_url": self.use_url,
        }
        if image_url:
            form["image_urls"] = [image_url]
        elif image_base64:
            form["binary_data_base64"] = [image_base64]
        resp = self.cv_process(form)
        return self.__common_handler_response(resp)

    def image_creator(self, image_url: str = None, image_base64: str = None, prompt: str = "",
                      width: int = 512, height: int = 512, steps: int = 80, scale: float = 3.5,
                      negative_prompt: str = "Low quality", schedule_conf: str = "general_v20_9B_pe",
                      config: dict = None):
        config = config or {}
        form = {
            "req_schedule_conf": schedule_conf,
            "negative_prompt": negative_prompt,
            "scale": scale,
            "ddim_steps": steps,
            "width": width,
            "height": height,
            "return_url": self.use_url,
            "prompt": prompt
        }
        config.update(form)
        if image_url:
            config["image_urls"] = [image_url]
        elif image_base64:
            config["binary_data_base64"] = [image_base64]
        resp = self.cv_process(config)
        return self.__common_handler_response(resp)

    def pretty_to_face(self, image_url: str = None, image_base64: str = None, beauty_level: float = 1.0):
        form = {
            "req_key": "face_pretty",
            "do_risk": False,
            "multi_face": "1",
            "beauty_level": beauty_level,
        }
        if image_url:
            form["image_url"] = [image_url]
        elif image_base64:
            form["image_base64"] = [image_base64]
        resp = self.face_pretty(form)
        print("pretty_to_face", resp)
        if "code" in resp and resp["code"] == 10000:
            image = resp.get("data", {}).get("image", "")
            image_bytes = base64.b64decode(image)
            return OperateResult.success(data=[image_bytes])
        return OperateResult.warning(message="异常")

    def change_face(self, image_urls: List[str] = None, image_base64s: List[str] = None, face_type: str = "area",
                    source_similarity: float = 1.0,
                    config: dict = None):
        config = config or {}
        form = {
            "req_key": "face_swap3_6",
            "face_type": face_type,
            "source_similarity": source_similarity,
            "return_url": self.use_url,
        }
        config.update(form)
        print("change_face", config)
        if image_urls:
            config["image_urls"] = image_urls
        elif image_base64s:
            config["binary_data_base64"] = image_base64s
        resp = self.cv_process(config)
        print("change_face", resp)
        return self.__common_handler_response(resp)

    def faceswap_ai(self, image_urls: List[str] = None, image_base64s: List[str] = None, gpen: float = 0.9,
                    skin: float = 0.1,
                    config: dict = None):
        config = config or {}
        form = {
            "req_key": "faceswap_ai",
            "do_risk": False,
            "gpen": gpen,
            "skin": skin,
            "return_url": True
        }
        config.update(form)
        print("faceswap_ai", config)
        if image_urls:
            config["image_urls"] = image_urls
        elif image_base64s:
            config["binary_data_base64"] = image_base64s
        resp = self.cv_process(config)
        print("faceswap_ai", resp)
        return self.__common_handler_response(resp)

    def keep_face_20(self, image_urls: List[str] = None, image_base64s: List[str] = None, prompt: str = "",
                     scale: float = 3.5,
                     steps: int = 90, width: int = 512, height: int = 512, cfg_rescale: float = 0.7,
                     ref_ip_weight: float = 0.7, ref_id_weight: float = 0.36,
                     config: dict = None):
        config = config or {}
        form = {
            "req_key": "high_aes_ip_v20",
            "prompt": prompt,
            "desc_pushback": True,
            "seed": -1,
            "scale": scale,
            "ddim_steps": steps,
            "width": width,
            "height": height,
            "cfg_rescale": cfg_rescale,
            "ref_ip_weight": ref_ip_weight,
            "ref_id_weight": ref_id_weight,
            "use_sr": True,
            "return_url": self.use_url,
        }
        config.update(form)
        if image_urls:
            config["image_urls"] = image_urls
        elif image_base64s:
            config["binary_data_base64"] = image_base64s
        resp = self.cv_process(config)
        return self.__common_handler_response(resp)

    def keep_face_14(self, image_urls: List[str] = None, image_base64s: List[str] = None, prompt: str = "",
                     scale: float = 3.5,
                     steps: int = 90, width: int = 512, height: int = 512,
                     ref_ip_weight: float = 0.7, ref_id_weight: float = 0.36,
                     config: dict = None):
        config = config or {}
        form = {
            "req_key": "high_aes_general_v14_ip_keep",
            "prompt": prompt,
            "model_version": "general_v1.4_ip",
            "subject_prompt": "",
            "ref_ip_weight": ref_ip_weight,
            "ref_id_weight": ref_id_weight,
            "negative_prompt": "negative prompt：nsfw, nude, blurry, watermark, identifying mark, low resolution, mutated, lack of hierarchy",
            "seed": -1,
            "scale": scale,
            "ddim_steps": steps,
            "width": width,
            "height": height,
            "use_rephraser": False,
            "use_predict_tags": True,
            "use_sr": True,
            "sr_seed": -1,
            "sr_strength": 0.4,
            "sr_scale": 3.5,
            "sr_steps": 10,
            "return_url": self.use_url,
        }
        config.update(form)
        if image_urls:
            config["image_urls"] = image_urls
        elif image_base64s:
            config["binary_data_base64"] = image_base64s
        resp = self.cv_process(config)
        return self.__common_handler_response(resp)
